package demo.generator;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.baomidou.mybatisplus.annotation.DbType;
import com.baomidou.mybatisplus.annotation.FieldFill;
import com.baomidou.mybatisplus.generator.AutoGenerator;
import com.baomidou.mybatisplus.generator.config.DataSourceConfig;
import com.baomidou.mybatisplus.generator.config.IDbQuery;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.VelocityTemplateEngine;

public class MyAutoGenerator extends AutoGenerator {

	private static final Logger logger = LoggerFactory.getLogger(MyAutoGenerator.class);

	@Override
	public void execute() {
		logger.debug("==========================准备生成文件...==========================");
		// 初始化配置
		if (null == config) {
			config = new ConfigBuilder(this.getPackageInfo(), this.getDataSource(), this.getStrategy(),
					this.getTemplate(), this.getGlobalConfig());
			if (null != injectionConfig) {
				injectionConfig.setConfig(config);
			}
		}
		if (null == this.getTemplateEngine()) {
			// 为了兼容之前逻辑，采用 Velocity 引擎 【 默认 】
			this.setTemplateEngine(new VelocityTemplateEngine());
			// templateEngine = new VelocityTemplateEngine();
		}
		// 将所有表的每个字段都追加fill值为UPDATE
		allTableFieldsFill();
		// 模板引擎初始化执行文件输出
		this.getTemplateEngine().init(this.pretreatmentConfigBuilder(config)).mkdirs().batchOutput().open();
		logger.debug("==========================文件生成完成！！！==========================");

	}

	/**
	 * 将所有表的每个字段都追加fill值为UPDATE
	 */
	private void allTableFieldsFill() {
		ConfigBuilder configBuilder = this.getConfig();
		DataSourceConfig dataSourceConfig = this.getDataSource();
		IDbQuery dbQuery = dataSourceConfig.getDbQuery();
		try (Connection connection = dataSourceConfig.getConn()) {
			List<TableInfo> tableInfoList = configBuilder.getTableInfoList();
			for (TableInfo tableInfo : tableInfoList) {
				oneTableFieldsFill(tableInfo, dbQuery, connection, dataSourceConfig);
			}
		} catch (Exception e) {
			System.err.println("SQL Exception：" + e.getMessage());
		}

	}

	/**
	 * 将一个表的每个字段都追加fill值为UPDATE
	 * @param tableInfo
	 * @param dbQuery
	 * @param connection
	 * @param dataSourceConfig
	 */
	private void oneTableFieldsFill(TableInfo tableInfo, IDbQuery dbQuery, Connection connection,
			DataSourceConfig dataSourceConfig) {
		// 模仿ConfigBuilder.convertTableFields来实现业务逻辑
		Map<String, TableField> tableFieldMap = getTableFieldMap(tableInfo);
		String tableFieldsSql = getTableFieldsSql(tableInfo, dbQuery, dataSourceConfig);
		try (PreparedStatement preparedStatement = connection.prepareStatement(tableFieldsSql);
				ResultSet results = preparedStatement.executeQuery()) {
			List<TableField> tableFieldList = new ArrayList<>();
			while (results.next()) {
				String fieldName = results.getString(dbQuery.fieldName());
				TableField tableField = tableFieldMap.get(fieldName);
				if (tableField != null) {
					tableField.setFill(FieldFill.UPDATE.toString());
					tableFieldList.add(tableField);
				}
			}
			tableInfo.setFields(tableFieldList);
		} catch (Exception e) {
			System.err.println("SQL Exception：" + e.getMessage());
		}
	}

	/**
	 * 取得字段名-字段对象映射
	 * @param tableInfo
	 * @return
	 */
	private Map<String, TableField> getTableFieldMap(TableInfo tableInfo) {
		return tableInfo.getFields().stream().collect(Collectors.toMap(TableField::getName, Function.identity()));
	}

	/**
	 * 获取查询表字段SQL
	 * @param tableInfo
	 * @param dbQuery
	 * @param dataSourceConfig
	 * @return
	 */
	private String getTableFieldsSql(TableInfo tableInfo, IDbQuery dbQuery, DataSourceConfig dataSourceConfig) {
		String tableFieldsSql = dbQuery.tableFieldsSql();
		if (DbType.POSTGRE_SQL == dbQuery.dbType()) {
			tableFieldsSql = String.format(tableFieldsSql, dataSourceConfig.getSchemaName(), tableInfo.getName());
		} else if (DbType.ORACLE == dbQuery.dbType()) {
			tableFieldsSql = String.format(tableFieldsSql.replace("#schema", dataSourceConfig.getSchemaName()),
					tableInfo.getName());
		} else {
			tableFieldsSql = String.format(tableFieldsSql, tableInfo.getName());
		}
		return tableFieldsSql;
	}
}
